#!/usr/bin/env python
# -*- coding: UTF-8 -*-
'''
@Project ：V2 
@File    ：T.py
@IDE     ：PyCharm 
@Author  ：郭星
@Date    ：2025/10/7 20:16 
'''
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Paper Experiment Reproduction (Python, single-file)
[...docstring trimmed for brevity in this saved file...]
"""

import math
import random
import argparse
from dataclasses import dataclass
from typing import List, Tuple, Dict, Any
import numpy as np
import matplotlib.pyplot as plt

N_LINES = 3
N_MACH  = 9
N_BUFF  = N_MACH - 1

MTBF = np.array([
    [280,270,282,288,296,270,334,318,248],
    [264,294,262,280,279,310,284,302,268],
    [278,285,326,295,269,284,317,297,260],
], dtype=float)

MTTR_min = np.array([
    [28,26,28,28.5,29,26.5,32,27,24],
    [26,28.5,27,28,27,30,28,29.5,25],
    [28,27,30,27.5,28,30,31,28.5,27],
], dtype=float)

CAP_BOUNDS = np.array([
    [[15,20],[8,13],[16,23],[8,14],[13,21],[17,21],[12,15],[9,13],[10,22]],
    [[15,18],[15,22],[9,12],[15,23],[8,13],[12,15],[17,23],[10,20],[9,18]],
    [[17,22],[12,16],[15,18],[12,13],[14,23],[10,14],[8,12],[14,18],[10,20]],
], dtype=float)

TRANSFER_COST = np.array([
    [0,5,0,5,0,5,6,0],
    [0,5,4,5,4,5,0,0],
    [0,0,4,0,4,0,6,0],
], dtype=float)

REVENUE = {'A':4.0,'B':4.5,'C':4.0,'D':5.0,'E':3.5}
AVG_REVENUE = float(np.mean(list(REVENUE.values())))

def availability(MTBF_hours: float, MTTR_minute: float) -> float:
    MTTR_hours = MTTR_minute / 60.0
    return MTBF_hours / (MTBF_hours + MTTR_hours + 1e-12)

def sa_ratio(w_up: float, w_down: float) -> float:
    return w_up / (w_up + w_down + 1e-12)

def pv_has_vacancy(Sa: float, C: int) -> float:
    eps = 1e-12
    Sa = max(Sa, eps)
    if abs(Sa - 1.0) < 1e-9:
        return C / (C + 1.0)
    num = (Sa**(C+1) - Sa)
    den = (Sa**(C+1) - 1.0)
    return max(0.0, min(1.0, num/(den + eps)))

def pt_has_inventory(Sa: float, C: int) -> float:
    return max(0.0, min(1.0, Sa * pv_has_vacancy(Sa, C)))

def ps_block_prob(Sa: float, C: int) -> float:
    Sa = max(0.0, min(1.0, Sa))
    return max(0.0, min(1.0, Sa * (1.0 / (1.0 + 0.15*C))))

def pu_idle_prob(Sa: float, C_prev: int) -> float:
    Sa = max(0.0, min(1.0, Sa))
    base = max(0.0, 1.0 - 2.0*Sa)
    damp = 1.0 / (1.0 + 0.1 * C_prev)
    return max(0.0, min(1.0, base * damp))

def w_matrix_from_bounds(bounds: np.ndarray) -> np.ndarray:
    return bounds.mean(axis=-1)

@dataclass
class Candidate:
    C: np.ndarray
    Ts: np.ndarray
    W: float = 0.0
    theta: float = 0.0
    details: Dict[str, Any] = None

def legalize(C: np.ndarray, Ts: np.ndarray, C_max: int):
    C = np.maximum(0, np.round(C).astype(int))
    tot = C.sum()
    if tot > C_max:
        scale = C_max / (tot + 1e-12)
        C = np.floor(C * scale).astype(int)
    mask_shareable = (TRANSFER_COST > 0).astype(float)
    Ts = np.clip(Ts, 0.0, 1.0) * mask_shareable
    return C, Ts

def evaluate(C: np.ndarray, Ts: np.ndarray):
    W_nom = w_matrix_from_bounds(CAP_BOUNDS)
    A = np.zeros_like(W_nom)
    for i in range(N_LINES):
        for j in range(N_MACH):
            A[i,j] = availability(MTBF[i,j], MTTR_min[i,j])
    W_eff = W_nom * A

    line_outputs = []
    sharing_cost = 0.0

    for i in range(N_LINES):
        stage_rates = []
        for j in range(N_BUFF):
            w_up   = W_eff[i,j]
            w_down = W_eff[i,j+1]
            Sa     = sa_ratio(w_up, w_down)

            P_v = pv_has_vacancy(Sa, int(C[i,j]))
            P_t = pt_has_inventory(Sa, int(C[i,j]))
            P_s = ps_block_prob(Sa, int(C[i,j]))
            P_u = pu_idle_prob(Sa, int(C[i,j]))

            ts = Ts[i,j]
            shareable = TRANSFER_COST[i,j] > 0
            if shareable and ts > 0:
                P_v_eff = min(1.0, P_v + 0.25 * ts * (1.0 - P_v))
                P_t_eff = min(1.0, P_t + 0.25 * ts * (1.0 - P_t))
            else:
                P_v_eff = P_v
                P_t_eff = P_t

            Rb_eff = (1.0 - P_u - P_s) * P_v_eff * w_up
            Rm_eff = (1.0 - P_u - P_s) * P_t_eff * w_down

            stage_rate = min(Rb_eff, Rm_eff)
            stage_rates.append(stage_rate)

            sharing_cost += ts * float(C[i,j]) * TRANSFER_COST[i,j]

        line_outputs.append(min(stage_rates))

    line_outputs = np.array(line_outputs)
    W_total = AVG_REVENUE * line_outputs.sum() - sharing_cost

    ideal_outputs = []
    for i in range(N_LINES):
        stage_nom = []
        for j in range(N_BUFF):
            w_up   = W_nom[i,j]
            w_down = W_nom[i,j+1]
            stage_nom.append(min(w_up, w_down))
        ideal_outputs.append(min(stage_nom))
    ideal_outputs = np.array(ideal_outputs)
    theta = float((line_outputs / (ideal_outputs + 1e-9)).mean())
    details = dict(line_outputs=line_outputs, sharing_cost=sharing_cost, ideal_outputs=ideal_outputs)
    return float(W_total), theta, details

def non_dominated_sort(points: List[tuple]) -> List[List[int]]:
    N = len(points)
    S = [set() for _ in range(N)]
    n = [0]*N
    fronts = [[]]
    for p in range(N):
        for q in range(N):
            if p==q: continue
            if points[p][0] >= points[q][0] and points[p][1] >= points[q][1] \
               and (points[p][0] > points[q][0] or points[p][1] > points[q][1]):
                S[p].add(q)
            elif points[q][0] >= points[p][0] and points[q][1] >= points[p][1] \
               and (points[q][0] > points[p][0] or points[q][1] > points[p][1]):
                n[p] += 1
        if n[p]==0:
            fronts[0].append(p)
    i = 0
    while fronts[i]:
        Q = []
        for p in fronts[i]:
            for q in S[p]:
                n[q] -= 1
                if n[q]==0:
                    Q.append(q)
        i += 1
        fronts.append(Q)
    if not fronts[-1]:
        fronts.pop()
    return fronts

def crowding_distance(front_points: List[tuple]) -> List[float]:
    m = len(front_points)
    if m==0:
        return []
    dist = [0.0]*m
    for dim in range(2):
        idx = sorted(range(m), key=lambda k: front_points[k][dim])
        dist[idx[0]] = dist[idx[-1]] = float('inf')
        vmin = front_points[idx[0]][dim]
        vmax = front_points[idx[-1]][dim]
        denom = vmax - vmin + 1e-12
        for r in range(1, m-1):
            dist[idx[r]] += (front_points[idx[r+1]][dim] - front_points[idx[r-1]][dim]) / denom
    return dist

def firefly_mo_search(C_max=1500, pop=40, iters=25, seed=42):
    rng = np.random.default_rng(seed)

    def init_C_random():
        return rng.integers(0, 50, size=(N_LINES, N_BUFF))

    def init_C_heur_gauss():
        C = np.zeros((N_LINES, N_BUFF), dtype=int)
        for i in range(N_LINES):
            x = np.arange(N_BUFF)
            mu = (N_BUFF-1)/2.0
            sigma = 1.8
            w = np.exp(-0.5*((x-mu)/sigma)**2)
            w = w / (w.sum()+1e-12)
            base = int(C_max / N_LINES)
            Ci = np.floor(w * base).astype(int)
            C[i,:] = Ci
        return C

    def init_C_prod_based():
        C = np.zeros((N_LINES, N_BUFF), dtype=int)
        W_nom = w_matrix_from_bounds(CAP_BOUNDS)
        for i in range(N_LINES):
            w = W_nom[i,:-1]
            w = w / (w.sum()+1e-12)
            base = int(C_max / N_LINES)
            Ci = np.floor(w * base).astype(int)
            C[i,:] = Ci
        return C

    def init_Ts_from_pretransfer(C):
        W_eff = w_matrix_from_bounds(CAP_BOUNDS) * np.vectorize(availability)(MTBF, MTTR_min/60.0)
        Ts = np.zeros((N_LINES, N_BUFF), dtype=float)
        for i in range(N_LINES):
            for j in range(N_BUFF):
                if TRANSFER_COST[i,j] <= 0:
                    Ts[i,j] = 0.0
                else:
                    w_up = W_eff[i,j]
                    w_dn = W_eff[i,j+1]
                    Sa   = sa_ratio(w_up, w_dn)
                    Ts[i,j] = max(0.0, min(1.0, 0.2 + 0.6*(Sa-0.5)))
        return Ts

    population: List[Candidate] = []
    for k in range(pop):
        if k < pop*0.3:
            C0 = init_C_random()
        elif k < pop*0.6:
            C0 = init_C_heur_gauss()
        else:
            C0 = init_C_prod_based()
        Ts0 = init_Ts_from_pretransfer(C0)
        C0, Ts0 = legalize(C0, Ts0, C_max)
        W0, th0, det0 = evaluate(C0, Ts0)
        population.append(Candidate(C0, Ts0, W0, th0, det0))

    ATTR0 = 1.0
    GAMMA = 1.0
    SIGMA0 = 0.10
    SIGMAF = 0.03

    for it in range(iters):
        t = (it+1)/iters
        pts = [(c.W, c.theta) for c in population]
        fronts = non_dominated_sort(pts)

        front0 = fronts[0] if fronts else []
        crowd = {}
        if front0:
            f0_pts = [pts[idx] for idx in front0]
            cd = crowding_distance(f0_pts)
            for k, idx in enumerate(front0):
                crowd[idx] = cd[k]

        new_pop: List[Candidate] = []
        for idx_a in range(len(population)):
            a = population[idx_a]
            candidates = []
            for f, front in enumerate(fronts):
                for idx in front:
                    af = next((g for g,fr in enumerate(fronts) if idx_a in fr), 10**9)
                    if f < af:
                        candidates.append((f, idx))
                    elif f == af:
                        candidates.append((f - (crowd.get(idx,0.0)*1e-3), idx))
            if not candidates:
                b = random.choice(population)
            else:
                candidates.sort(key=lambda x: (x[0], -pts[x[1]][0], -pts[x[1]][1]))
                b = population[candidates[0][1]]

            # move
            de = np.linalg.norm((a.C - b.C).astype(float))
            dt = np.linalg.norm((a.Ts - b.Ts).astype(float))
            attr = ATTR0 * math.exp(-GAMMA*(de + dt))
            sigma = SIGMA0*(1.0 - t) + SIGMAF*t

            C_new = a.C.astype(float) + attr*(b.C - a.C) + np.random.normal(0, sigma, size=a.C.shape)
            Ts_new = a.Ts + attr*(b.Ts - a.Ts) + np.random.normal(0, sigma, size=a.Ts.shape)

            C_new, Ts_new = legalize(C_new, Ts_new, C_max)
            if random.random() < 0.4:
                i = random.randrange(N_LINES)
                j = random.randrange(N_BUFF)
                step = np.random.normal(0, sigma)
                Ts_new[i,j] = np.clip(Ts_new[i,j] + step, 0.0, 1.0 if TRANSFER_COST[i,j]>0 else 0.0)

            Wn, thn, detn = evaluate(C_new, Ts_new)
            new_pop.append(Candidate(C_new, Ts_new, Wn, thn, detn))

        combined = population + new_pop
        pts_comb = [(c.W, c.theta) for c in combined]
        fronts_comb = non_dominated_sort(pts_comb)

        population_next: List[Candidate] = []
        for front in fronts_comb:
            if len(population_next) + len(front) <= pop:
                population_next += [combined[idx] for idx in front]
            else:
                fpts = [pts_comb[idx] for idx in front]
                cd = crowding_distance(fpts)
                order = sorted(range(len(front)), key=lambda k: cd[k], reverse=True)
                remain = pop - len(population_next)
                population_next += [combined[front[k]] for k in order[:remain]]
                break

        population = population_next

    pts_final = [(c.W, c.theta) for c in population]
    fronts_final = non_dominated_sort(pts_final)
    pareto = [population[idx] for idx in fronts_final[0]] if fronts_final else population
    return pareto, population

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--cmax", type=int, default=1500)
    parser.add_argument("--pop", type=int, default=40)
    parser.add_argument("--iters", type=int, default=25)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--sweep", action="store_true")
    args = parser.parse_args()

    pareto, pop = firefly_mo_search(C_max=args.cmax, pop=args.pop, iters=args.iters, seed=args.seed)

    plt.figure(figsize=(6,4))
    Ws = [c.W for c in pop]; Th = [c.theta for c in pop]
    plt.scatter(Ws, Th, s=20, alpha=0.6, label="Population")
    Ws_p = [c.W for c in pareto]; Th_p = [c.theta for c in pareto]
    plt.scatter(Ws_p, Th_p, s=35, marker='x', label="Pareto (approx)")
    plt.xlabel("Average earnings proxy (W)")
    plt.ylabel("Throughput proxy (θ)")
    plt.legend()
    plt.tight_layout()
    plt.savefig("pareto_demo.png", dpi=180)
    print("Saved Pareto plot -> pareto_demo.png")

    best = max(pareto, key=lambda c: (c.W, c.theta))
    print("Best candidate summary:")
    print("  W =", best.W, "theta =", best.theta)
    print("  Line outputs (units/min):", best.details["line_outputs"])
    print("  Sharing cost proxy:", best.details["sharing_cost"])

    np.save("best_C.npy", best.C)
    np.save("best_Ts.npy", best.Ts)
    print("Saved best_C.npy, best_Ts.npy")

    if args.sweep:
        from collections import OrderedDict
        C_list = [300, 600, 900]
        summary = OrderedDict()
        for Cmax in C_list:
            pareto_i, _ = firefly_mo_search(C_max=Cmax, pop=max(24, args.pop//2), iters=max(15, args.iters//2), seed=args.seed+int(Cmax))
            best_i = max(pareto_i, key=lambda c: c.W + 1000*c.theta)
            summary[Cmax] = dict(W=best_i.W, theta=best_i.theta)
        import csv
        with open("cmax_sweep_summary.csv", "w", newline="") as f:
            w = csv.writer(f); w.writerow(["C_max","W_proxy","theta_proxy"])
            for Cmax, rec in summary.items():
                w.writerow([Cmax, rec["W"], rec["theta"]])
        print("Saved sweep summary -> cmax_sweep_summary.csv")

if __name__ == "__main__":
    main()
